/*
 *  See layout.h to get memory layout.
 *  1. Build a virtual memory map for kernel, the map is simple:
 *        linearaddr = phyaddr - 3G, 0 <= phyaddr < MIN(1G-4M, memory_size)
 *     The reason to keep this map simple is that: we can only get the
 *     physical address of users' page tables, by the page directory.
 *     We have to know their linear address to modify the PTE(page table
 *     entry), in order to handle the page fault interupt. We can get
 *     their linear address easily thanks to this simple map.
 *  2. As a result, the kernel cann't access the memory outside of 1G-4M.
 *     So, we keep the highest 4M for dynamic map. map_page() can be used
 *     to map any page to 4G-4M ~ 4G, so the kernel can access through the
 *     memory.
 *  3. The kernel's page tables begin from 1M, followed by mm's bitmap.
 *  4. We try to alloc pages at 0 ~ 1G, to kernel. So search free pages
 *     from bitmap's beginning for kernel, and from end for user's process.
 *
 */


#include <type.h>
#include "tty.h"
#include "process.h"
#include "layout.h"

#define PAGE_SHIFT	12
#define PAGE_SIZE	(1<<PAGE_SHIFT)

/* Attributes for page directory and page table. */
#define PF_P		0x001	// present
#define PF_RW		0x002	// read&write
#define PF_US		0x004	// user or super
#define PF_WT		0x008	// write through
#define PF_CD		0x010	// cache disable
#define PF_A		0x020	// accessed
#define PF_G		0x100	// global
/* Attributes only for page directory. */
#define PDF_S		0x080	// page size
/* Attributes only for page table. */
#define PTF_D		0x040	// dirty
#define PTF_AT		0x080	// page tabel attribute index

/* Page Entry Mask: address and attribute. */
#define PEM_ADDRESS	0xFFFFF000
#define PEM_ATTRIBUTE	0x00000FFF


static u8_t	*Bitmap_addr;
static int	Bitmap_size;
static int	Pages_free;


void init_mm(void)
{
/* See layout.h to get more details of the memory layout. */

	/* ==== Get memory's size. ==== */
	struct memory_info {
		u32_t	BaseAddrLow, BaseAddrHigh;
		u32_t	LengthLow, LengthHigh;
		u32_t	Type;
	} *mem_info = (struct memory_info *)INFO_MEMORY;
	int mem_info_num = *((int *)INFO_MEMORY_NUM);
	u32_t mem_size;	// size of memory, in byte
	int i;

	for(i = 0; i < mem_info_num; i++) {
		if(mem_info[i].Type == 1)
			mem_size = mem_info[i].BaseAddrLow
				+ mem_info[i].LengthLow;
	}
	mem_size = (mem_size + 0xFFFFF) & 0xFFF00000;
	printk("Memory(mega): %d\n", mem_size >> 20);


	/* ==== Set page tables for kernel. ==== */
	u32_t *page_dir = (u32_t *)KERNEL_PAGE_DIR;
	u32_t *page_table = (u32_t *)KERNEL_PAGE_TABLES;
	u32_t pde, pte;	// page directory/table entry
	int pde_num;	// number of pde
	int pt_size;	// size of page tables, in 32-bit

	/* Min(mem_size, 1G-4M). Keep the last pde(4M) for dynamic map. */
	pde_num = MIN(mem_size >> 22, 0xFF); // 22 : 4M each pde
	/* page tables. */
	pt_size = (pde_num - 1) << 10;
	pte = 0x0 | PF_P | PF_RW | PF_US; // TODO:no PF_US
	for(i = 0; i < pt_size; i++) {
		page_table[i] = pte;
		pte += 0x1000;
	}
	/* page directory. */
	pde = PHY_ADDR_PAGE_TABLES | PF_P | PF_RW | PF_US; // TODO:no PF_US
	for(i = 0; i < pde_num; i++) {
		page_dir[i + 0x300] = pde;
		pde += 0x1000;
	}
	page_dir[1023] = PHY_ADDR_PAGE_TABLE_HIGH | PF_P | PF_RW;
	memset(KERNEL_PAGE_TABLE_HIGH, 0, PAGE_SIZE);


	/* ==== Initilize page bitmap. ==== */
	u32_t kernel_pages, bitmap_pages, used_pages;

	/* Zero the bitmap. */
	//Bitmap_addr = (u8_t *)(PHY_ADDR_PAGE_TABLES + (pde_num << PAGE_SHIFT) + VIR_KERNEL);
	//Bitmap_size = mem_size >> 15;	// size in byte
	Bitmap_addr = (u8_t *)PHY2LIN(PHY_ADDR_PAGE_TABLES + (pde_num << PAGE_SHIFT));
	Bitmap_size = mem_size >> 12;	// size in byte
	memset(Bitmap_addr, 0, Bitmap_size);

	/* Mark the used memory: kernel(1M) + page-tables + bitmap. */
	kernel_pages = PHY_ADDR_PAGE_TABLES >> PAGE_SHIFT;
	bitmap_pages = (Bitmap_size + 0xFFF) >> PAGE_SHIFT;
	used_pages = kernel_pages + pde_num + bitmap_pages;
	for(i = 0; i < used_pages; i++) {
		Bitmap_addr[i] = 85; // none zero
	}

	/* Init pages in free.*/
	Pages_free = (mem_size >> PAGE_SHIFT) - used_pages;
}

static u32_t alloc_page(int flag)
{
	static u32_t last = 0;
	u32_t i;

	if(Pages_free == 0)
		return 0;

	for(i = last; i < Bitmap_size; i++) {
		if(Bitmap_addr[i] == 0)
			goto out;
	}
	for(i = 0; i < last; i++) {
		if(Bitmap_addr[i] == 0)
			goto out;
	}
out:
	Pages_free--;
	Bitmap_addr[i] = 1;
	last = i;
	return i;
}
static void free_page(u32_t page_nr)
{
	Bitmap_addr[page_nr]--;
	if(Bitmap_addr[page_nr] == 0)
		Pages_free++;
}

/* Map @phy_addr to @vir_addr. If @phy_addr=0, allocate
 * a page, and return its physical address.  */
u32_t set_page_table(u32_t vir_addr, u32_t phy_addr)
{
	int i_dir, i_table;
	u32_t pde;
	u32_t *page_dir, *page_table;

	page_dir = (u32_t *)PHY2LIN(current->tss.cr3_pdbr);
	i_dir = vir_addr >> 22;
	if((page_dir[i_dir] & 1) == 0) {
		/* Page table does not exist too. */
		pde = alloc_page(0) << PAGE_SHIFT;
		page_dir[i_dir] = pde | PF_P | PF_RW | PF_US;
		page_table = (u32_t *)PHY2LIN(pde);
		memset(page_table, 0, 0x1000);
	} else {
		page_table = (u32_t *)PHY2LIN(page_dir[i_dir] & PEM_ADDRESS);
	}

	i_table = (vir_addr >> PAGE_SHIFT) & 0x3FF;
	if(phy_addr == 0)
		phy_addr = alloc_page(0) << PAGE_SHIFT;
	else if(page_table[i_table] != 0)
		; // we need do something here
	page_table[i_table] = phy_addr | PF_P | PF_RW | PF_US;

	return phy_addr;
}

/* Map @phy_addr to 0xFFC00000, so the kernel can access the memory
 * whose address is larger than 1G.
 * TODO: 0xFFC00000 - 4G. */
u32_t map_page(u32_t phy_addr)
{
	u32_t *page_table = (u32_t *)KERNEL_PAGE_TABLE_HIGH;
	page_table[0] = phy_addr  | PF_P | PF_RW;
	return 0xFFC00000;
}

void unmap_page(u32_t lin_addr)
{
	u32_t *page_table = (u32_t *)KERNEL_PAGE_TABLE_HIGH;
	page_table[0] = 0;
}

static int do_page_wp(u32_t cr2)
{
	u32_t *pdir, *ptable1, *ptable2;
	u32_t lin_addr, phy_addr;
	int i_dir, i_table;
	int imap_table, imap_page;

	pdir = (u32_t *)PHY2LIN(current->tss.cr3_pdbr);
	i_dir = cr2 >> 22;
	ptable1 = PHY2LIN(pdir[i_dir] & PEM_ADDRESS);
	i_table = cr2 << 10 >> 22;
	imap_page = ptable1[i_table] >> PAGE_SHIFT;

	/* If the page is shared by more than one process, copy
	 * the page and mark its PTE read-write. Else, just mark PTE. */
	if(Bitmap_addr[imap_page] > 1) {

		/* Same with the page, check the page table is shared or not. */
		imap_table = pdir[i_dir] >> PAGE_SHIFT;
		if(Bitmap_addr[imap_table] > 1) {
			phy_addr = alloc_page(0) << 12;
			ptable2 = PHY2LIN(phy_addr);
			memcpy(ptable2, ptable1, PAGE_SIZE);
			pdir[i_dir] = phy_addr | (pdir[i_dir] & PEM_ATTRIBUTE);

			Bitmap_addr[imap_table]--;

			/* Invalid TLB entry. */
			flush_tlb_entry(ptable1);
		} else {
			ptable2 = ptable1;
		}

		/* Alloc and copy physical page. */
		phy_addr = alloc_page(0) << 12;
		lin_addr = map_page(phy_addr);
		memcpy(lin_addr, cr2 & 0xFFFFF000, PAGE_SIZE);
		unmap_page(lin_addr);
		ptable2[i_table] =  phy_addr | (ptable2[i_table] & PEM_ATTRIBUTE) | PF_RW;

		Bitmap_addr[imap_page]--;

		/* Invalid TLB entry. */
		flush_tlb_entry(cr2 & 0xFFFFF000);
	} else {
		ptable1[i_table] |= PF_RW;
	}
}

void do_page_fault(u32_t error_code, u32_t cr2)
{
	if((error_code & PF_P) == 0) {	/* Not present. */
		set_page_table(cr2, 0);
	} else {			/* Write protect.(error_code & PF_RW) */
		do_page_wp(cr2);
	}
}

/* @flag_pl: return physical of linear address. */
u32_t new_page(int flag_pl)
{
	u32_t addr = alloc_page(0) << PAGE_SHIFT;
	return flag_pl==0 ? PHY2LIN(addr) : addr;
}

void increase_page(u32_t addr)
{
	addr &= 0x3FFFFFFF;
	Bitmap_addr[addr >> PAGE_SHIFT]++;
}

void delete_page(u32_t addr)
{
	addr &= 0x3FFFFFFF;
	free_page(addr >> PAGE_SHIFT);
}

/* Do something related with page table when fork.
 * Only copy page directory, and mark PDE and PTE, 
 * so the father and child process share page tables
 * and physical pages. */
u32_t page_table_fork(void)
{
	u32_t *pdir, *ptable;
	u32_t pdir_new;
	int i, j;

	/* Mark all existing PTE read-only. */
	pdir = (u32_t *)PHY2LIN(current->tss.cr3_pdbr);
	for(i = 1; i < 768; i++) { /* XXX: should from 0. */
		if(pdir[i] == 0)
			continue;
		ptable = (u32_t *)PHY2LIN(pdir[i] & PEM_ADDRESS);
		Bitmap_addr[pdir[i] >> 12]++;
		for(j = 0; j < 1024; j++) {
			if(ptable[j] == 0)
				continue;
			ptable[j] &= ~PF_RW;
			Bitmap_addr[ptable[j] >> 12]++;
		}
	}

	/* Flush TLB. */
	flush_tlb();

	/* Copy page directory. */
	pdir_new = new_page(1);
	memcpy(PHY2LIN(pdir_new), pdir, 4096);

	return pdir_new;
}

/* @flag_ro: 0 - free all pages
 *           1 - only free shared pages */
void free_page_table(int flag_ro)
{
	u32_t *pdir, *ptable;
	int i, j;
	int remain;

	pdir = (u32_t *)PHY2LIN(current->tss.cr3_pdbr);
	for(i = 1; i < 768; i++) { /* XXX: should from 0 */
		if(pdir[i] == 0)
			continue;
		ptable = (u32_t *)PHY2LIN(pdir[i] & PEM_ADDRESS);

		/* Free all pages, or the page-table is shared (which
		 * all the pages belonging to it are read-only),
		 * then free the pages and page-table, and clear PDE.
		 * But don't clear the PTE! */
		if(flag_ro == 0 || Bitmap_addr[pdir[i]>>12] > 1) {
			for(j = 0; j < 1024; j++) {
				if(ptable[j] != 0)
					delete_page(ptable[j]);
			}
			delete_page(pdir[i]);
			pdir[i] = 0;
		}

		/* The page table is not shared, then only free the
		 * shared pages and clear their PTEs. If all pages are
		 * shared, free the page-table and clear PDE. */
		else {
			remain = 0;
			for(j = 0; j < 1024; j++) {
				if(ptable[j] == 0)
					continue;

				if(Bitmap_addr[ptable[j] >> 12] > 1) {
					delete_page(ptable[j]);
					ptable[j] = 0;
				} else {
					remain = 1;
				}
			}
			if(remain == 0) {
				delete_page(pdir[i]);
				pdir[i] = 0;
			}
		}
	}
}

void do_debug(void)
{
	u32_t *pdir, *ptable;
	int i, j;

	u32_t vir_addr = 0xB0000000;

	i = vir_addr >> 22;
	j = (vir_addr >> PAGE_SHIFT) & 0x3FF;
	pdir = (u32_t *)PHY2LIN(current->tss.cr3_pdbr);
	ptable = (u32_t *)PHY2LIN(pdir[i] & PEM_ADDRESS);
	printk("(%d) vir-phy:  %d,%d: %x-%x\n", current->pid, i, j, vir_addr & PEM_ADDRESS, ptable[j]);
}
